"""
数据库工具函数
"""
import logging
from typing import List, Dict, Any, Tuple, Type
from sqlalchemy import insert, select, update, and_
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.dialects.mysql import insert as mysql_insert

logger = logging.getLogger(__name__)

def update_or_insert(db, model: Type, data_list: List[Dict[str, Any]], unique_fields: List[str]) -> Tuple[int, int]:
    """
    通用的更新或插入函数，支持MySQL和SQLite
    
    Args:
        db: SQLAlchemy会话
        model: 数据库模型类
        data_list: 要插入或更新的数据列表，每项是一个字典
        unique_fields: 用于判断记录是否存在的唯一字段列表
        
    Returns:
        tuple: (inserted_count, updated_count) 插入和更新的记录数
    """
    if not data_list:
        return 0, 0
        
    inserted_count = 0
    updated_count = 0
    
    try:
        for data in data_list:
            # 构建查询条件，检查记录是否存在
            conditions = []
            for field in unique_fields:
                conditions.append(getattr(model, field) == data[field])
            
            # 查询是否存在匹配的记录
            query = select(model).where(and_(*conditions))
            existing_record = db.execute(query).scalar_one_or_none()
            
            if existing_record:
                # 记录存在，执行更新
                # 排除唯一字段，只更新其他字段
                update_data = {k: v for k, v in data.items() if k not in unique_fields and k != 'id'}
                if update_data:  # 如果有需要更新的字段
                    stmt = update(model).where(and_(*conditions)).values(**update_data)
                    db.execute(stmt)
                    updated_count += 1
            else:
                # 记录不存在，执行插入
                stmt = insert(model).values(**data)
                db.execute(stmt)
                inserted_count += 1
        
        db.commit()
        return inserted_count, updated_count
        
    except SQLAlchemyError as e:
        db.rollback()
        logger.error(f"数据库操作失败: {e}", exc_info=True)
        raise

def batch_update_or_insert(db, model: Type, data_list: List[Dict[str, Any]], unique_fields: List[str], 
                           batch_size: int = 1000) -> Tuple[int, int]:
    """
    批量更新或插入函数，支持处理大量数据
    
    Args:
        db: SQLAlchemy会话
        model: 数据库模型类
        data_list: 要插入或更新的数据列表，每项是一个字典
        unique_fields: 用于判断记录是否存在的唯一字段列表
        batch_size: 每批处理的记录数，默认1000
        
    Returns:
        tuple: (inserted_count, updated_count) 插入和更新的记录数
    """
    total_inserted = 0
    total_updated = 0
    
    # 按批次处理数据
    for i in range(0, len(data_list), batch_size):
        batch = data_list[i:i+batch_size]
        inserted, updated = update_or_insert(db, model, batch, unique_fields)
        total_inserted += inserted
        total_updated += updated
        
    return total_inserted, total_updated 